import os
# 使用国内镜像
os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
from transformers import BertTokenizer, BertModel
import torch
from sklearn.metrics.pairwise import cosine_similarity

# 加载预训练的 BERT 模型和分词器
tokenizer = BertTokenizer.from_pretrained('bert-base-chinese')
model = BertModel.from_pretrained('bert-base-chinese')

# 示例句子（取自《荷塘月色》）
sentence1 = "月光下的荷塘，曲曲折折的荷塘上，弥望的是田田的叶子。"
sentence2 = "月光如流水一般，静静地泻在这一片叶子和花上。"

# 生成句子的词嵌入
def generate_embeddings(sentence, tokenizer, model):
    inputs = tokenizer(sentence, return_tensors="pt", truncation=True, max_length=128)
    outputs = model(**inputs)
    # 提取最后一层的隐藏状态作为嵌入向量
    embeddings = outputs.last_hidden_state
    return embeddings, inputs['input_ids']

# 获取两个句子的嵌入向量
embeddings1, input_ids1 = generate_embeddings(sentence1, tokenizer, model)
embeddings2, input_ids2 = generate_embeddings(sentence2, tokenizer, model)

# 比较"月"在两种上下文中的向量差异
def get_word_embedding(word, embeddings, input_ids, tokenizer):
    token_index = torch.where(input_ids[0] == tokenizer.convert_tokens_to_ids(word))[0][0]
    return embeddings[0][token_index].detach().numpy()

# 提取"月"在句子中的嵌入
word = "月"
word_embedding1 = get_word_embedding(word, embeddings1, input_ids1, tokenizer)
word_embedding2 = get_word_embedding(word, embeddings2, input_ids2, tokenizer)

# 计算余弦相似度
similarity = cosine_similarity([word_embedding1], [word_embedding2])[0][0]
print(f"'月' 在两个上下文中的向量相似性: {similarity:.4f}")

# 打印向量信息
print("\n'月' 在句子 1 中的嵌入向量 (前5维):", word_embedding1[:5])
print("'月' 在句子 2 中的嵌入向量 (前5维):", word_embedding2[:5])